package icu.zhhll.wrapper;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * 解决Response只可读取一次的问题
 *
 * @author zh
 * @date 2023/8/14 17:37
 */
public class ResponseReaderHttpServletResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream output;
    private ServletOutputStream filterOutput;

    public ResponseReaderHttpServletResponseWrapper(HttpServletResponse response) {
        super(response);
        this.output = new ByteArrayOutputStream();
    }

    public ServletOutputStream getOutputStream() throws IOException {
        if (this.filterOutput == null) {
            this.filterOutput = new ServletOutputStream() {

                @Override
                public void write(int b) throws IOException {
                    output.write(b);
                }

                @Override
                public boolean isReady() {
                    return false;
                }

                @Override
                public void setWriteListener(WriteListener writeListener) {

                }
            };
        }
        return this.filterOutput;
    }
}
